import sys
import csv
import pickle
import lib.tools as my_tools

from tqdm import tqdm

def sort_dataset(): 
    data_set = my_tools.load_dataset('found_data_set.pickle')   

    # 对数据集进行排序，并保存到csv文件中
    print(f"Waiting for sorting dataset")
    sorted_dataset = []
    with open('dataset.csv', 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        # 写入表头
        writer.writerow(['D_value', 'D1', 'D2', 'D3', 'v_zd', 'omega_x', 'omega_y', 'omega_z', 'p_x', 'p_y', 'c_z', 'gv_z', 'n3'])
    
        for i in tqdm(range(len(data_set)), total = len(data_set)):
            
            if len(data_set[i]) > 0:
                try:
                    data_set[i].sort(key=lambda x: x[0])
                except Exception as e:
                    print(f"Error while sorting data_set[{i}]: {e}")
                    print(f"data_set[{i}] is {data_set[i]}")
                writer.writerow(data_set[i][0])
                sorted_dataset.append(data_set[i][0])
        # 同时将筛选后的数据直接保存在pickle文件中，便于后续训练使用
        my_tools.save_dataset(sorted_dataset, 'sorted_dataset.pickle')

def sort_dataset_n3():
   
    # data_set = my_tools.load_dataset('found_data_set_attitude.pickle')

    eta = 0.5
    # 对数据集进行排序，并保存到csv文件中
    print(f"Waiting for sorting dataset")
    sorted_dataset_attitude = []
    with open('dataset_attitude.csv', 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        # 写入表头
        writer.writerow(['D_value', 'omega_x', 'omega_y', 'omega_z', 'n3_x', 'n3_y', 'n3_z', 'n3_xd', 'n3_yd', 'n3_zd','D+eta × norm2_x^2'])

        # chunk_count = my_tools.count_chunks('found_data_set_attitude.pickle')
        chunk_count = 32
        dataset_reader = my_tools.load_dataset_chunk('found_data_set_attitude.pickle')
        for data_set in tqdm(dataset_reader, desc="Processing data chunks"):
            for i in range(len(data_set)):
                
                if len(data_set[i]) > 0:
                    # data_set[i] 列表中的第一列数据D-value, 加上特定值形如\eta||x||^2, 附加到每行末尾
                    for j in range(len(data_set[i])):
                        data_set[i][j].append(abs(data_set[i][j][0] + eta*(data_set[i][j][4] - data_set[i][j][7])**2 + eta*(data_set[i][j][5] - data_set[i][j][8])**2)) 
                    # 对数据集进行排序
                    try:
                        data_set[i].sort(key=lambda x: x[10])
                    except Exception as e:
                        print(f"Error while sorting data_set[{i}]: {e}")
                        print(f"data_set[{i}] is {data_set[i]}")
                    writer.writerow(data_set[i][0])
                    sorted_dataset_attitude.append(data_set[i][0])
        # 同时将筛选后的数据直接保存在pickle文件中，便于后续训练使用
        my_tools.save_dataset(sorted_dataset_attitude, 'sorted_dataset_attitude.pickle')

def sort_dataset_pg():
   
    data_set = my_tools.load_dataset('found_data_set_pg.pickle')

    # 对数据集进行排序，并保存到csv文件中
    print(f"Waiting for sorting dataset")
    sorted_dataset_pg = []
    with open('dataset_pg.csv', 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        # 写入表头
        writer.writerow(['D_value', 'v_zd', 'omega_x', 'omega_y', 'omega_z', 'p_x', 'p_y', 'c_z', 'gv_z', 'n3'])
    
        for i in tqdm(range(len(data_set)), total = len(data_set)):
            
            if len(data_set[i]) > 0:
                try:
                    data_set[i].sort(key=lambda x: x[0])
                except Exception as e:
                    print(f"Error while sorting data_set[{i}]: {e}")
                    print(f"data_set[{i}] is {data_set[i]}")
                writer.writerow(data_set[i][0])
                sorted_dataset_pg.append(data_set[i][0])
        # 同时将筛选后的数据直接保存在pickle文件中，便于后续训练使用
        my_tools.save_dataset(sorted_dataset_pg, 'sorted_dataset_pg.pickle')


COMMANDS = {
    "all": "Sort full dataset",
    "attitude": "Sort attitude control dataset",
    "pg": "Sort p_g control dataset"
}

if __name__ == "__main__":
    if len(sys.argv) == 1:
        my_tools.print_help(sys.argv[0], COMMANDS)

    if len(sys.argv) == 2:
        if sys.argv[1] == 'all':
            sort_dataset()
        elif sys.argv[1] == 'attitude':
            sort_dataset_n3()
        else:
            my_tools.print_help(sys.argv[0], COMMANDS)